package com.xinggq.pay.service.config;

import com.xinggq.common.context.TraceContext;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.TaskDecorator;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.Map;
import java.util.concurrent.Executor;
import java.util.concurrent.ThreadPoolExecutor;

/**
 * 异步配置类
 * 配置异步任务执行器，支持TraceId传递和线程名优化
 * 注意，任何使用@Async注解的方法，都应该指定配置在这个类下的线程池，以避免TraceID为空的情况
 *
 * @author xinggq
 * @since 2025/1/10
 */
@Slf4j
@Configuration
@EnableAsync
public class AsyncConfig {

    @Value("${async.executor.core-pool-size:5}")
    private int corePoolSize;

    @Value("${async.executor.max-pool-size:20}")
    private int maxPoolSize;

    @Value("${async.executor.queue-capacity:100}")
    private int queueCapacity;

    @Value("${async.executor.thread-name-prefix:async-}")
    private String threadNamePrefix;

    @Value("${async.executor.keep-alive-seconds:60}")
    private int keepAliveSeconds;

    /**
     * 异步任务执行器
     * 配置线程池参数并支持TraceId传递
     */
    @Bean("taskExecutor")
    public Executor taskExecutor() {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        
        // 线程池配置
        executor.setCorePoolSize(corePoolSize);
        executor.setMaxPoolSize(maxPoolSize);
        executor.setQueueCapacity(queueCapacity);
        executor.setThreadNamePrefix(threadNamePrefix);
        executor.setKeepAliveSeconds(keepAliveSeconds);
        
        // 拒绝策略：由调用线程执行
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
        
        // 设置任务装饰器，用于传递TraceId和MDC
        executor.setTaskDecorator(new TraceTaskDecorator());
        
        // 等待所有任务完成后再关闭线程池
        executor.setWaitForTasksToCompleteOnShutdown(true);
        executor.setAwaitTerminationSeconds(60);
        
        executor.initialize();
        
        log.info("异步线程池初始化完成 - 核心线程数: {}, 最大线程数: {}, 队列容量: {}", 
                corePoolSize, maxPoolSize, queueCapacity);
        
        return executor;
    }

    /**
     * 任务装饰器，用于传递TraceId和MDC上下文，并设置带方法名的线程名
     */
    private static class TraceTaskDecorator implements TaskDecorator {
        
        private static final String PROJECT_PACKAGE = "com.xinggq.pay.service";
        private static final String[] EXCLUDED_METHODS = {
            "getAsyncMethodName", "decorate", "lambda$", "run", "call"
        };
        
        @Override
        public Runnable decorate(Runnable runnable) {
            // 获取当前线程的上下文信息
            String traceId = TraceContext.getTraceId();
            Map<String, String> contextMap = MDC.getCopyOfContextMap();
            String methodName = getAsyncMethodName();
            
            return () -> {
                String originalThreadName = Thread.currentThread().getName();
                try {
                    // 设置增强的线程名
                    setEnhancedThreadName(originalThreadName, methodName);
                    
                    // 设置上下文
                    setThreadContext(traceId, contextMap);
                    
                    // 执行原始任务
                    runnable.run();
                } finally {
                    // 清理资源
                    cleanupThreadContext(originalThreadName);
                }
            };
        }
        
        /**
         * 设置增强的线程名
         */
        private void setEnhancedThreadName(String originalName, String methodName) {
            if (methodName != null) {
                Thread.currentThread().setName(originalName + "-" + methodName);
            }
        }
        
        /**
         * 设置线程上下文
         */
        private void setThreadContext(String traceId, Map<String, String> contextMap) {
            if (traceId != null) {
                TraceContext.setTraceId(traceId);
            }
            if (contextMap != null) {
                MDC.setContextMap(contextMap);
            }
        }
        
        /**
         * 清理线程上下文
         */
        private void cleanupThreadContext(String originalThreadName) {
            Thread.currentThread().setName(originalThreadName);
            TraceContext.clear();
            MDC.clear();
        }
        
        /**
         * 获取异步方法名
         */
        private String getAsyncMethodName() {
            try {
                StackTraceElement[] stackTrace = Thread.currentThread().getStackTrace();
                
                for (StackTraceElement element : stackTrace) {
                    String className = element.getClassName();
                    String methodName = element.getMethodName();
                    
                    // 检查是否为项目包下的方法且不在排除列表中
                    if (isProjectMethod(className) && !isExcludedMethod(methodName)) {
                        return formatMethodName(className, methodName);
                    }
                }
            } catch (Exception e) {
                log.debug("获取异步方法名失败: {}", e.getMessage());
            }
            return null;
        }
        
        /**
         * 检查是否为项目方法
         */
        private boolean isProjectMethod(String className) {
            return className.startsWith(PROJECT_PACKAGE);
        }
        
        /**
         * 检查是否为排除的方法
         */
        private boolean isExcludedMethod(String methodName) {
            for (String excluded : EXCLUDED_METHODS) {
                if (methodName.equals(excluded) || methodName.startsWith(excluded)) {
                    return true;
                }
            }
            return false;
        }
        
        /**
         * 格式化方法名
         */
        private String formatMethodName(String className, String methodName) {
            String simpleClassName = className.substring(className.lastIndexOf('.') + 1);
            return simpleClassName + "." + methodName;
        }
    }
}